{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RI points location\n",
    "dusurui     2023072212\n",
    "gaemi       2024071912\n",
    "haikui      2023090112\n",
    "kangni      2024102812\n",
    "koinu       2023100112\n",
    "mojie       2024090300\n",
    "saola       2023082600\n",
    "shantuo     2024092800"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.lines as mlines\n",
    "import matplotlib.dates as mdates\n",
    "import cartopy.crs as ccrs\n",
    "from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter\n",
    "import cartopy.feature as cfeature\n",
    "import shapely.geometry as sgeom\n",
    "from datetime import datetime,timedelta\n",
    "import tqdm\n",
    "from global_land_mask import globe\n",
    "from typlot.scripts.gsj_typhoon import tydat,see,count_rapidgrow,split_str_id,tydat_CMA, basemap\n",
    "from geopy.distance import geodesic\n",
    "import cartopy.io.shapereader as shpreader\n",
    "import geopandas as gpd\n",
    "from typlot.config.global_config import *\n",
    "\n",
    "\n",
    "initimes = [\"2023072212\",\"2023090112\",\"2023082600\",\"2023100112\",\"2024090300\",\"2024071912\",\"2024102812\",\"2024092800\"]\n",
    "\n",
    "tynames,tyids = split_str_id(names)\n",
    "\n",
    "# 🔧 改进1: 调整figsize和间距参数\n",
    "fig = plt.figure(figsize=(10, 14))  # 增加整体尺寸\n",
    "gs = fig.add_gridspec(4, 2, \n",
    "                      hspace=0.20,  # 垂直间距\n",
    "                      wspace=0.20,  # 水平间距\n",
    "                      left=0.05,    # 左边距\n",
    "                      right=0.95,   # 右边距\n",
    "                      top=0.96,     # 顶部边距\n",
    "                      bottom=0.04)  # 底部边距\n",
    "\n",
    "axes = []\n",
    "for i in range(4):\n",
    "    for j in range(2):\n",
    "        ax = fig.add_subplot(gs[i, j], projection=ccrs.PlateCarree())\n",
    "        axes.append(ax)\n",
    "\n",
    "for axi,(ty,tyid,initime) in enumerate(zip(tynames,tyids,initimes)):\n",
    "    ax = axes[axi]\n",
    "    RIstd = 7\n",
    "    dt = timedelta(hours=48)  \n",
    "    directory = os.path.join(global_ensdir, f\"{ty}_{tyid}\")\n",
    "    dates_name = [initime]\n",
    "    dates_paths = [os.path.join(directory,date) for date in dates_name]\n",
    "    obs_path = os.path.join(global_obsdir, f'{ty}_CMAobs.txt')\n",
    "    pic_savepath = paperdir \n",
    "    os.makedirs(pic_savepath,exist_ok=True)\n",
    "    \n",
    "    def track(ty, ax, lon, lat, title=None, color=None, marker=\".\", label=None, show_label=False, s=12, time_data=None):\n",
    "        \"\"\"绘制台风轨迹（时间分段显示）\"\"\"\n",
    "        if time_data is not None:\n",
    "            time_array = np.array(time_data)\n",
    "        else:\n",
    "            time_array = ty.time\n",
    "        \n",
    "        now = start\n",
    "        i = 0\n",
    "        colors = plt.cm.tab20(np.linspace(0, 1, 20))\n",
    "        \n",
    "        lon = np.array(lon)\n",
    "        lat = np.array(lat)\n",
    "        \n",
    "        if len(time_array) > 0:\n",
    "            total_duration = max(time_array) - start\n",
    "            max_iterations = int(total_duration / dt) + 2\n",
    "        else:\n",
    "            max_iterations = 20\n",
    "        \n",
    "        iteration = 0\n",
    "        while now <= max(time_array) and iteration < max_iterations:\n",
    "            end = now + dt\n",
    "            \n",
    "            if time_data is not None:\n",
    "                want = (time_array >= now) & (time_array < end)\n",
    "                if hasattr(ty, 'time') and len(ty.time) > 0:\n",
    "                    try:\n",
    "                        record_mask = (record == 1)\n",
    "                        want = want & record_mask\n",
    "                    except:\n",
    "                        pass\n",
    "            else:\n",
    "                want = (time_array >= now) & (time_array < end)\n",
    "                try:\n",
    "                    want = want & (record == 1)\n",
    "                except:\n",
    "                    pass\n",
    "            \n",
    "            if show_label:\n",
    "                hours_from_start = int((now - start).total_seconds() / 3600)\n",
    "                hours_to_end = int((end - start).total_seconds() / 3600)\n",
    "                if dt.total_seconds() >= 24 * 3600:\n",
    "                    label_to_use = f\"{dt.days*i}-{dt.days*(i+1)}d\"\n",
    "                else:\n",
    "                    label_to_use = f\"{hours_from_start}h-{hours_to_end}h\"\n",
    "            else:\n",
    "                label_to_use = None\n",
    "            \n",
    "            if np.any(want):\n",
    "                ax.scatter(\n",
    "                    lon[want], lat[want], \n",
    "                    marker=marker, \n",
    "                    transform=ccrs.PlateCarree(), \n",
    "                    color=colors[i % len(colors)],\n",
    "                    label=label_to_use, \n",
    "                    s=s,\n",
    "                    alpha=0.8,\n",
    "                    edgecolors='white',\n",
    "                    linewidths=0.5,\n",
    "                    zorder=10  # 🔧 确保点在上层\n",
    "                )\n",
    "            \n",
    "            now = end\n",
    "            i += 1\n",
    "            iteration += 1\n",
    "\n",
    "    def track_of_members(ax, ty, time, lat, lon, marker='.', show_label=True, s=24):\n",
    "        \"\"\"绘制集合成员轨迹\"\"\"\n",
    "        now = start\n",
    "        i = 0\n",
    "        colors = plt.cm.tab20(np.linspace(0, 1, 20))\n",
    "        time = np.array(time)\n",
    "        lat = np.array(lat)\n",
    "        lon = np.array(lon)\n",
    "        \n",
    "        if len(time) > 0:\n",
    "            max_time = min(start + timedelta(days=16), max(time))\n",
    "        else:\n",
    "            max_time = start + timedelta(days=16)\n",
    "        \n",
    "        iteration = 0\n",
    "        max_iterations = int((max_time - start) / dt) + 2\n",
    "        \n",
    "        while now <= max_time and iteration < max_iterations:\n",
    "            end = now + dt\n",
    "            want = (time >= now) & (time < end)\n",
    "            \n",
    "            if show_label:\n",
    "                hours_from_start = int((now - start).total_seconds() / 3600)\n",
    "                hours_to_end = int((end - start).total_seconds() / 3600)\n",
    "                if dt.total_seconds() >= 24 * 3600:\n",
    "                    label_to_use = f\"{dt.days*i}-{dt.days*(i+1)}d\"\n",
    "                else:\n",
    "                    label_to_use = f\"{hours_from_start}h-{hours_to_end}h\"\n",
    "            else:\n",
    "                label_to_use = None\n",
    "            \n",
    "            if np.any(want):\n",
    "                ax.scatter(\n",
    "                    lon[want], lat[want], \n",
    "                    marker=marker, \n",
    "                    transform=ccrs.PlateCarree(), \n",
    "                    color=colors[i % len(colors)], \n",
    "                    label=label_to_use, \n",
    "                    s=s,\n",
    "                    alpha=0.7,\n",
    "                    edgecolors='white',\n",
    "                    linewidths=0.5,\n",
    "                    zorder=10\n",
    "                )\n",
    "            \n",
    "            now = end\n",
    "            i += 1\n",
    "            iteration += 1\n",
    "    \n",
    "    # 起报时间循环\n",
    "    for i in tqdm.tqdm(range(len(dates_paths)), desc=f\"Processing {ty}\"):\n",
    "        members_paths = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f != \"TRACK_ID_0\" and f.startswith(\"TRACK_ID\")]\n",
    "        era_path = [os.path.join(dates_paths[i], f) for f in os.listdir(dates_paths[i]) if f.startswith(\"TRACK_ID_0\")][0]\n",
    "        \n",
    "        start = datetime.strptime(dates_name[i], \"%Y%m%d%H\")\n",
    "        \n",
    "        print(f\"处理起报时间: {start}\")\n",
    "        \n",
    "        '''draw members'''\n",
    "        time = []\n",
    "        lat = []\n",
    "        lon = []\n",
    "        \n",
    "        for j in range(len(members_paths)):\n",
    "            try:\n",
    "                member = tydat(members_paths[j], RIstd, lon_format='180')\n",
    "                record = member.num_rapidgrow()\n",
    "                RI = (record == 1)\n",
    "                if RI.any():\n",
    "                    time.extend(member.time[RI].tolist())\n",
    "                    lat.extend(member.lat[RI].tolist())\n",
    "                    lon.extend(member.lon[RI].tolist())\n",
    "            except Exception as e:\n",
    "                print(f\"处理成员 {j} 时出错: {e}\")\n",
    "                continue\n",
    "        \n",
    "        print(f\"集合成员RI点数量: {len(time)}\")\n",
    "        \n",
    "        if len(time) > 0:\n",
    "            track_of_members(ax, member, time, lat, lon, marker=\"o\", show_label=True, s=64)\n",
    "        \n",
    "        ''' draw Analysis '''\n",
    "        try:\n",
    "            era = tydat(era_path, RIstd, lon_format='180')\n",
    "            record = era.num_rapidgrow()\n",
    "            track(era, ax, era.lon, era.lat, marker=\"s\", s=96, show_label=False, time_data=era.time)\n",
    "        except Exception as e:\n",
    "            print(f\"处理分析数据时出错: {e}\")\n",
    "        \n",
    "        ''' obs '''\n",
    "        try:\n",
    "            obs = tydat_CMA(obs_path, lon_format='180')\n",
    "            record = count_rapidgrow(RIstd, obs.umax, obs.time)\n",
    "            track(obs, ax, obs.lon, obs.lat, marker=\"^\", s=96, show_label=False, time_data=obs.time)\n",
    "        except Exception as e:\n",
    "            print(f\"处理观测数据时出错: {e}\")\n",
    "        \n",
    "        basemap(ax, extent=[105,135,10,30], lon_interval=2.5, lat_interval=2.5)\n",
    "        \n",
    "        # 🔧 改进2: 优化标题位置\n",
    "        # title = f\"{dates_name[i]} {ty}\"\n",
    "        # ax.set_title(title, fontsize=10, pad=10)\n",
    "\n",
    "\n",
    "        \n",
    "        # 🔧 改进3: 优化图例位置和样式\n",
    "        handles, labels = ax.get_legend_handles_labels()\n",
    "        if len(handles) > 0:\n",
    "            # 去重（保持顺序）\n",
    "            unique_labels = []\n",
    "            unique_handles = []\n",
    "            for handle, label in zip(handles, labels):\n",
    "                if label not in unique_labels:\n",
    "                    unique_labels.append(label)\n",
    "                    unique_handles.append(handle)\n",
    "            \n",
    "            # 图例放在右侧外部\n",
    "            ax.legend(unique_handles, unique_labels,\n",
    "                     loc='center left',\n",
    "                     bbox_to_anchor=(0, 0.8),\n",
    "                     fontsize=8,\n",
    "                     frameon=True,\n",
    "                     fancybox=True,\n",
    "                     shadow=True,\n",
    "                     ncol=1)\n",
    "        \n",
    "        # 🔧 改进4: 优化说明文字位置\n",
    "        ax.text(x=0.02, y=0.02, \n",
    "                s=\"○ Ensemble  □ Analysis  △ Obs\", \n",
    "                transform=ax.transAxes, \n",
    "                color=\"black\", \n",
    "                fontsize=8, \n",
    "                va='bottom',\n",
    "                bbox=dict(boxstyle=\"round,pad=0.4\", \n",
    "                         facecolor=\"white\", \n",
    "                         alpha=0.85,\n",
    "                         edgecolor='gray'))\n",
    "\n",
    "\n",
    "    ax.text(\n",
    "        0.63, 0.06,\n",
    "        f\"{ty}({initimes[axi]})\",\n",
    "        transform=ax.transAxes,\n",
    "        fontsize=8,\n",
    "        fontweight='bold',\n",
    "        color='white',\n",
    "        verticalalignment='bottom',\n",
    "        horizontalalignment='left',\n",
    "        bbox=dict(\n",
    "            boxstyle='round,pad=0.5',\n",
    "            edgecolor='none',\n",
    "            alpha=0.7\n",
    "        )\n",
    "    )\n",
    "\n",
    "# 🔧 改进5: 使用 tight_layout 但保留图例空间\n",
    "\n",
    "fig.suptitle('RI Locations', fontsize=15, fontweight='bold',y=0.999)\n",
    "\n",
    "plt.tight_layout(rect=[0, 0, 0.98, 1])  # 右侧留空间给图例\n",
    "\n",
    "\n",
    "# 在循环结束后添加标签\n",
    "labels = [\"(a)\",\"(b)\",\"(c)\",\"(d)\",\"(e)\",\"(f)\",\"(g)\",\"(h)\"]\n",
    "for ax, label in zip(axes, labels):\n",
    "    ax.text(0.88, 0.92, label,\n",
    "            transform=ax.transAxes,  # 使用轴坐标系 (0-1)\n",
    "            fontsize=14,\n",
    "            fontweight='bold',\n",
    "            va='top',                # 垂直对齐：顶部\n",
    "            ha='left',              # 水平对齐：右侧\n",
    "            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "if draw_opt == False: \n",
    "    plt.savefig( os.path.join(paperdir,f'fig5.{pictype}'), \n",
    "                dpi=300, bbox_inches='tight')\n",
    "    plt.close()\n",
    "else:\n",
    "    plt.show()\n",
    "\n",
    "print(\"完成所有子图\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wmq",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
